package com.zzyl;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ChatCompletionClient {

    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;
    private final String apiUrl = "https://qianfan.baidubce.com/v2/chat/completions";
    private final String authorizationToken;

    // 构造方法，接收认证令牌
    public ChatCompletionClient(String authorizationToken) {
        this.restTemplate = new RestTemplate();
        this.objectMapper = new ObjectMapper();
        this.authorizationToken = authorizationToken;
    }

    // 发送聊天请求并返回原始响应
    public String sendChatRequest(String userMessage) {
        try {
            // 创建请求头
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("Authorization", authorizationToken);

            // 构建请求体
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "ernie-3.5-8k");

            List<Map<String, String>> messages = new ArrayList<>();

            Map<String, String> systemMessage = new HashMap<>();
            systemMessage.put("role", "system");
            systemMessage.put("content", "平台助手");
            messages.add(systemMessage);

            Map<String, String> userMsg = new HashMap<>();
            userMsg.put("role", "user");
            userMsg.put("content", userMessage);
            messages.add(userMsg);

            requestBody.put("messages", messages);

            // 创建HTTP实体
            HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);

            // 发送POST请求
            ResponseEntity<String> response = restTemplate.postForEntity(
                    apiUrl,
                    requestEntity,
                    String.class
            );

            // 处理响应
            if (response.getStatusCode().is2xxSuccessful()) {
                return response.getBody();
            } else {
                throw new RuntimeException("API请求失败: " + response.getStatusCode() +
                        ", 响应体: " + response.getBody());
            }
        } catch (Exception e) {
            throw new RuntimeException("发送请求失败: " + e.getMessage(), e);
        }
    }

    // 提取响应中的content内容（修正了解析路径）
    public String extractContent(String jsonResponse) {
        try {
            // 解析JSON响应
            JsonNode rootNode = objectMapper.readTree(jsonResponse);

            // 百度千帆API的正确响应结构路径
            return rootNode
                    .get("choices")
                    .get(0)
                    .get("message")
                    .get("content")
                    .asText();
        } catch (Exception e) {
            throw new RuntimeException("解析响应内容失败: " + e.getMessage() +
                    ", 原始响应: " + jsonResponse, e);
        }
    }

    // 主方法，用于测试
    public static void main(String[] args) {
        // 请替换为实际有效的授权令牌
        String token = "Bearer bce-v3/ALTAK-TF5a2JoOV8WqrW0SsuB5G/df096e19fed0ba26df21ae97a3c13f5fd5cc8213";
        ChatCompletionClient client = new ChatCompletionClient(token);

        try {
            // 发送请求获取原始响应
            String rawResponse = client.sendChatRequest("你好");
            System.out.println("原始响应: " + rawResponse);

            // 提取content内容
            String content = client.extractContent(rawResponse);
            System.out.println("提取的内容: " + content);
        } catch (Exception e) {
            System.err.println("发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
    